%--------------------------------------------------------------------------------
% AMTH 358 Final - Question 1
% Radix-2, decimation-in-time, FFT and IFFT
% Forward DFT: xf(k) = Sum{n=1,N-1} x(k) * exp(-i*2*pi*n*k/N)
% Inverse DFT: x(k) = Sum{n=1,N-1} (1/N) * xf(k) * exp(i*2*pi*n*k/N)
% Note that the scaling factor 1/N is applied to the inverse FFT mode
% following same behaviour as Matlab's native fft() and ifft() 
% Calls to fft_r2dt with invalid arguments return a vector with all
% its elements set to NaN 
%
% arg1 x       : input vector 
% arg2 inverse : set to 1 to perform ifft               (default is 0)
% arg3 debug   : set to 1 to dump more console messages (default is 1)
% arg4 strict  : set to 1 to enforce vector size limit  (default is 1)
%--------------------------------------------------------------------------------

function [xf] = fft_r2dt(x, inverse, debug, strict)
  if nargin <4
    strict=1;   % enforce 1024 input vector size limit
  end
  if nargin <3
    debug=1;    % print out debug and info messages to console
  end
  if nargin <2 
    inverse=0;  % forward FFT is the default
  end
  if nargin <1  % print help usage message to console
    usage();
    xf = [];
    return;
  end
  
  N=length(x);
  L=fix(log2(N));
  xf = NaN * ones(1,N); % output vector is initialized in case of error condition    

  if N ~= 2^L
    debug_disp('Error: input vector x must have size as power of 2',debug);
    return;
  elseif L>10
    if strict
      debug_disp('Error: input vector length must have a size of 1024 or less',debug);
      return;    
    else
      debug_disp('Info: input vector is pretty big!',debug);
    end
  elseif L==0  % must cover this corner case too!
    debug_disp('Info: input vector has length 1', debug);
    xf=x;
    return;
  end

  debug_disp('Info: N roots of unity:',debug);

  W=gen_unit_roots(N,inverse);
  debug_disp(W,debug);
  
  % plot_roots(W);  % for debugging
  
  debug_disp('Info: Input data vector in bit-reversed-address sequence',debug);

  xf = bit_rev_sort(x);
  debug_disp(xf,debug);
  
  % construct the radix-2 decimation in time butterfly network
  for stage=0:L-1
    xf_next = [];
    
    k1 = 1;
    W_index = 1;
    
    for bfly=0:N/2 -1
      k2 = k1 + 2^stage;
      
      din = [xf(k1) xf(k2)];
      dout = butterfly(din, W(W_index));

      xf_next(k1) = dout(1);
      xf_next(k2) = dout(2);
      
      if mod(bfly+1,2^stage) == 0
        k1=k1+ (2^stage)+1;
        W_index = 1;
      else
        k1=k1+1;
        W_index = W_index + 2^(L-1-stage);
      end
    end
    xf = xf_next;
  end

  % scaling by 1/N factor done on the inverse transform
  if inverse
    xf = xf/N;
  end
  
function [dout] = butterfly(din, coefficient)
  din2_scaled = din(2) * coefficient;
  dout(1) = din(1) + din2_scaled;
  dout(2) = din(1) - din2_scaled;
  
function [x_br] = bit_rev_sort(x)
  N=length(x);
  L=fix(log2(N));
  x_br = zeros(1,N);
  for i=1:N
    binary_index_minus1 = dec2bin(i-1,L);
    bit_rev_index = bin2dec(fliplr(binary_index_minus1))+1;
    x_br(bit_rev_index) = x(i);
  end
  
  
function [W] = gen_unit_roots(N,inverse)
  % generate a table containing Nth roots of unity
  % for inverse FFT, generate the conjugate roots
  W = [];
  for k=0:N-1
    root = exp(-i*2*pi*k/N);
    if inverse
      root=conj(root);
    end
    W = [W root];
  end
  
function plot_roots(W)
  % plot the N roots of unity to gain intuition
  figure;
  plot(W,'*');
  grid;
  title('N unity roots');
  xlabel('Real'); ylabel('Imag');
  hold on;
  syms k real;
  kr = linspace(0,2*pi,128);
  c = exp(i*k);
  cr = subs(c,k,kr);
  plot(cr,'-.');
  hold off;

function debug_disp(str, debug)
  if debug == 1
    disp(str);
  end

function usage()
  disp('Usage: fft_dit <data vector> [debug] [strict]');
  disp('               data vector has length equal to a power of two');
  disp('               debug = 1    print out messages to console (default)');
  disp('               strict = 1   enforce limits on vector size (default)');
  
  